import sys
sys.path.append('..')
import os
import glob
import pandas as pd
import numpy as np
import json
import argparse
from tqdm import tqdm
from utils.helper import get_subdirs
from pathlib import Path
from PIL import Image

def parse_args():
    """
    Set args parameters
    """
    parser = argparse.ArgumentParser(description='Split data for training Seg model.')

    parser.add_argument("--data_path", type=str, default='/opt/data/private/project/T5/1E853/train_data/Images', help="Define the ratio of validation set.")
    parser.add_argument("--label_path", type=str, default='/opt/data/private/project/T5/1E853/data_yolo_pred/Annotations_of_pred', help="Define the ratio of validation set.")
    parser.add_argument("--no_defect_code", nargs='+', default=['TNNNF0'], help="无缺陷类")
    args = parser.parse_args()

    return args


# 统计站点图像数量(含标注，无缺陷code除外)
def main():
    args = parse_args()

    codes = os.listdir(args.data_path)
    memo = {}
    for code in codes:
        if not os.path.isdir(os.path.join(args.data_path, code)):
            continue
        imgs = os.listdir(os.path.join(args.data_path, code))
        if len(imgs) == 0:  # code对应图像文件不存在
            print('empty: ', os.path.join(args.data_path, code))
            continue
        if code in args.no_defect_code:  # 无缺陷code，无需考虑label
            memo[code] = len(imgs)
        elif not os.path.exists(os.path.join(args.label_path, code)):  # 有缺陷code，但是对应label文件夹不存在
            print('empty: ', os.path.join(args.data_path, code))
        else:  # 有缺陷code，统计label数量
            labels = os.listdir(os.path.join(args.data_path, code))
            memo[code] = img_label_intersection(imgs, labels)  # 取图像和label交集

    print('code num: ', len(memo.keys()))
    print('sum: ', sum(memo.values()))
    memo = sorted(memo.items(), key=lambda x:x[1], reverse=True)
    for key, value in memo:
        print('%s: %d'%(key, value))



def img_label_intersection(imgs, labels):
    imgs = [str(Path(x).stem) for x in imgs]
    labels = [str(Path(x).stem) for x in labels]

    return len(np.intersect1d(imgs, labels))


if __name__ == '__main__':
    main()
